package com.wxl.asmgradleplugin.myplugin

import com.android.build.api.transform.Format
import com.android.build.api.transform.QualifiedContent.ContentType
import com.android.build.api.transform.QualifiedContent.Scope
import com.android.build.api.transform.Transform
import com.android.build.api.transform.TransformInvocation
import javassist.ClassPool
import org.gradle.api.Project

import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassWriter
import org.objectweb.asm.Opcodes
import java.io.File

class MyTransform(private val project: Project) : Transform() {

    companion object {
        private const val DOT_CLASS = ".class"
    }

    override fun getName(): String = "MyTransform"

    override fun getInputTypes(): MutableSet<ContentType> = TransformManager.CONTENT_CLASS

    override fun getScopes(): MutableSet<in Scope> = TransformManager.SCOPE_FULL_PROJECT

    override fun isIncremental(): Boolean = false

    override fun transform(transformInvocation: TransformInvocation) {
        val inputs = transformInvocation.inputs
        val outputProvider = transformInvocation.outputProvider
        val classPool = ClassPool.getDefault()
        inputs.forEach { input ->
            input.directoryInputs.forEach { dirInput ->
                val inputDir = dirInput.file
                classPool.insertClassPath(inputDir.absolutePath)
                inputDir.walkTopDown().forEach { file ->
                    if (file.isFile && file.name.endsWith(DOT_CLASS)) {
                        val className = file.relativeTo(inputDir)
                            .path.removeSuffix(DOT_CLASS)
                            .replace(File.separator, ".")

                        if (className.startsWith(Consts.PK)) {
                            try {
                                println("====111 Transforming: $className")
//                                val ctClass = classPool.getCtClass(className)
//                                if (!ctClass.isFrozen) {
//                                    val method = ctClass.getDeclaredMethod("test")
//                                    method.insertBefore("""println(">>> Injected by Javassist <<<")""")
//                                    ctClass.writeFile(inputDir.absolutePath)
//                                    ctClass.detach()
//                                }
                                val classReader = ClassReader(file.readBytes())
                                val classWriter = ClassWriter(classReader, ClassWriter.COMPUTE_MAXS)
                                val classVisitor = LogClassVisitor(Opcodes.ASM9, classWriter)
                                classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES)
                                val bytes = classWriter.toByteArray()
                                file.writeBytes(bytes)
                            } catch (e: Exception) {
                                println("====111 Transform error: $className -> $e")
                            }
                        }
                    }
                }

                // 复制修改后的 class 到输出目录
                val outputDir = outputProvider.getContentLocation(
                    dirInput.name, dirInput.contentTypes, dirInput.scopes, Format.DIRECTORY
                )
                inputDir.copyRecursively(outputDir, overwrite = true)
            }
        }
    }
}
